# -*- coding: UTF-8 -*-
# __auth__ = "Liujue"
# pip install -r requirements.txt

import _thread as thread
import base64
import datetime
import hashlib
import hmac
import json
from urllib.parse import urlparse
import ssl
from datetime import datetime
from time import mktime
from urllib.parse import urlencode
from wsgiref.handlers import format_date_time
import websocket
from flask import Flask, request, jsonify, render_template
from flask_cors import CORS
import threading
import os
from dotenv import load_dotenv

# 加载环境变量
load_dotenv()

app = Flask(__name__)
CORS(app)

# 配置验证
required_env_vars = [
    'SPARK_APPID', 'SPARK_API_KEY', 'SPARK_API_SECRET',
    'SPARK_URL_X1', 'SPARK_URL_LITE',
    'SPARK_DOMAIN_X1', 'SPARK_DOMAIN_LITE'
]

for var in required_env_vars:
    if not os.getenv(var):
        raise ValueError(f"缺少必要的环境变量: {var}")

# 模型配置
MODEL_CONFIG = {
    "x1": {
        "appid": os.getenv('SPARK_APPID'),
        "api_key": os.getenv('SPARK_API_KEY'),
        "api_secret": os.getenv('SPARK_API_SECRET'),
        "spark_url": os.getenv('SPARK_URL_X1'),
        "domain": os.getenv('SPARK_DOMAIN_X1')
    },
    "lite": {
        "appid": os.getenv('SPARK_APPID'),
        "api_key": os.getenv('SPARK_API_KEY'),
        "api_secret": os.getenv('SPARK_API_SECRET'),
        "spark_url": os.getenv('SPARK_URL_LITE'),
        "domain": os.getenv('SPARK_DOMAIN_LITE')
    }
}


class WsParam:
    """WebSocket参数处理类"""

    def __init__(self, app_id, api_key, api_secret, spark_url):
        self.app_id = app_id
        self.api_key = api_key
        self.api_secret = api_secret
        self.host = urlparse(spark_url).netloc
        self.path = urlparse(spark_url).path
        self.spark_url = spark_url

    def create_url(self):
        """生成WebSocket连接URL"""
        now = datetime.now()
        date = format_date_time(mktime(now.timetuple()))

        signature_origin = f"host: {self.host}\ndate: {date}\nGET {self.path} HTTP/1.1"
        signature_sha = hmac.new(
            self.api_secret.encode('utf-8'),
            signature_origin.encode('utf-8'),
            digestmod=hashlib.sha256
        ).digest()
        signature_sha_base64 = base64.b64encode(signature_sha).decode('utf-8')

        authorization_origin = (
            f'api_key="{self.api_key}", algorithm="hmac-sha256", '
            f'headers="host date request-line", signature="{signature_sha_base64}"'
        )
        authorization = base64.b64encode(authorization_origin.encode('utf-8')).decode('utf-8')

        return f"{self.spark_url}?{urlencode({
            'authorization': authorization,
            'date': date,
            'host': self.host
        })}"


class ConversationManager:
    """对话历史管理类"""

    def __init__(self):
        self.conversations = {}

    def add_message(self, role, content, conversation_id):
        """添加消息到对话历史"""
        if conversation_id not in self.conversations:
            self.conversations[conversation_id] = []

        self.conversations[conversation_id].append({
            "role": role,
            "content": content,
            "timestamp": datetime.now().isoformat()
        })
        return self.conversations[conversation_id]

    def get_history(self, conversation_id):
        """获取对话历史"""
        return self.conversations.get(conversation_id, [])

    def get_conversation_length(self, conversation_id):
        """计算对话内容的长度"""
        history = self.get_history(conversation_id)
        return sum(len(msg["content"]) for msg in history)

    def trim_conversation(self, conversation_id, max_length=8000):
        """修剪对话历史，保持最大长度"""
        while self.get_conversation_length(conversation_id) > max_length:
            self.conversations[conversation_id].pop(0)
        return self.get_history(conversation_id)


conversation_manager = ConversationManager()


def on_message(ws, message, callback=None):
    """WebSocket消息处理"""
    data = json.loads(message)
    code = data['header']['code']

    if code != 0:
        error_msg = f'请求错误: {code}, {data}'
        if callback:
            callback({'error': error_msg, 'finished': True})
        ws.close()
        return

    choices = data["payload"]["choices"]
    status = choices["status"]
    text = choices['text'][0]

    response = {
        'finished': status == 2,
        'status': status
    }

    if 'reasoning_content' in text and text['reasoning_content']:
        response['reasoning_content'] = text["reasoning_content"]

    if 'content' in text and text['content']:
        response['content'] = text["content"]

    if callback:
        callback(response)

    if status == 2:
        ws.close()


def on_error(ws, error, callback=None):
    """WebSocket错误处理"""
    error_msg = f"WebSocket错误: {error}"
    if callback:
        callback({'error': error_msg, 'finished': True})


def on_close(ws, close_status_code, close_msg, callback=None):
    """WebSocket关闭处理"""
    if callback:
        callback({'finished': True, 'closed': True})


def on_open(ws):
    """WebSocket连接建立处理"""
    thread.start_new_thread(ws.run_thread, ())


def gen_params(appid, domain, messages):
    """生成请求参数"""
    return {
        "header": {"app_id": appid, "uid": "1234"},
        "parameter": {
            "chat": {
                "domain": domain,
                "temperature": 1.2,
                "max_tokens": 32768
            }
        },
        "payload": {"message": {"text": messages}}
    }


class SparkAPIClient:
    """星火API客户端"""

    def __init__(self, conversation_manager):
        self.conversation_manager = conversation_manager

    def ask_spark(self, config, user_message, conversation_id, callback=None):
        """调用星火API"""
        ws_param = WsParam(
            config["appid"],
            config["api_key"],
            config["api_secret"],
            config["spark_url"]
        )

        # 添加用户消息并修剪历史
        self.conversation_manager.add_message("user", user_message, conversation_id)
        messages = self.conversation_manager.trim_conversation(conversation_id)

        def message_handler(response):
            """处理API响应"""
            if 'content' in response:
                self.conversation_manager.add_message(
                    "assistant",
                    response['content'],
                    conversation_id
                )
            if callback:
                callback(response)

        ws = websocket.WebSocketApp(
            ws_param.create_url(),
            on_message=lambda ws, msg: on_message(ws, msg, message_handler),
            on_error=lambda ws, error: on_error(ws, error, message_handler),
            on_close=lambda ws, code, msg: on_close(ws, code, msg, message_handler),
            on_open=on_open
        )

        ws.appid = config["appid"]
        ws.domain = config["domain"]
        ws.question = messages
        ws.run_thread = lambda: ws.send(json.dumps(gen_params(
            ws.appid,
            ws.domain,
            ws.question
        )))

        # 启动WebSocket连接
        wst = threading.Thread(
            target=ws.run_forever,
            kwargs={"sslopt": {"cert_reqs": ssl.CERT_NONE}}
        )
        wst.daemon = True
        wst.start()

        return wst


spark_client = SparkAPIClient(conversation_manager)


@app.route('/')
def index():
    """首页"""
    return render_template('index.html')


@app.route('/api/chat', methods=['POST'])
def chat():
    """聊天API"""
    data = request.json
    user_message = data.get('message')
    conversation_id = data.get('conversation_id', f"conv_{datetime.now().timestamp()}")
    model_type = data.get('model_type', 'lite')

    if not user_message:
        return jsonify({"error": "消息不能为空"}), 400

    config = MODEL_CONFIG.get(model_type)
    if not config:
        return jsonify({"error": "无效的模型类型"}), 400

    responses = []
    spark_client.ask_spark(
        config,
        user_message,
        conversation_id,
        lambda msg: responses.append(msg)
    )

    return jsonify({
        "status": "processing",
        "conversation_id": conversation_id,
        "model_type": model_type
    })


@app.route('/api/history/<conversation_id>', methods=['GET'])
def get_history(conversation_id):
    """获取对话历史"""
    history = conversation_manager.get_history(conversation_id)
    return jsonify({"history": history})


@app.route('/api/models', methods=['GET'])
def get_models():
    """获取可用模型列表"""
    return jsonify({
        "models": [
            {"id": "lite", "name": "星火Lite模型"},
            {"id": "x1", "name": "星火X1模型"}
        ]
    })


if __name__ == '__main__':
    template_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'templates')
    app.template_folder = template_dir

    debug_mode = os.getenv('DEBUG', 'False').lower() == 'true'
    app.run(
        host="0.0.0.0",
        port=5000,
        debug=debug_mode
    )